{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Script for adversarial attacks with GNNs (including GIB-GAT, GAT, GCN and other baselines.)\"\"\"\n",
    "import argparse\n",
    "from copy import deepcopy\n",
    "import matplotlib.pylab as plt\n",
    "import numpy as np\n",
    "import os.path as osp\n",
    "import pickle\n",
    "from scipy import sparse\n",
    "import datetime\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import sys, os\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "from GIB.experiments.GIB_node_model import GNN, load_model_dict_GNN, train, get_data, edge_index_2_csr, get_attacked_data_deeprobust\n",
    "from GIB.pytorch_net.util import to_np_array, to_Variable, filter_filename, record_data, Beta_Function, Early_Stopping, make_dir, str2bool, to_string, eval_tuple, get_list_elements, plot_matrices\n",
    "from GIB.DeepRobust.deeprobust.graph.defense import GCN\n",
    "from GIB.DeepRobust.deeprobust.graph.defense import GCNJaccard\n",
    "from GIB.DeepRobust.deeprobust.graph.defense import RGCN\n",
    "from deeprobust.graph.data import Dataset\n",
    "from GIB.util import process_data_for_nettack, select_nodes, add_distant_neighbors, GIB_PATH\n",
    "import pdb\n",
    "\n",
    "is_cuda = \"cuda:0\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "METRICS_RECORD = [\"train_acc\", \"val_acc\", \"test_acc\", \"epochs\",\n",
    "                  \"train_loss\", \"val_loss\", \"test_loss\",\n",
    "                  \"train_f1_micro\", \"best_val_f1_micro\", \"b_test_f1_micro\",\n",
    "                  \"train_f1_macro\", \"val_f1_macro\", \"test_f1_macro\",\n",
    "                 ]\n",
    "\n",
    "def train_baseline_multiple(data, model_type, device, retrain_iters=5, suffix=\"\", verbose=True):\n",
    "    models = []\n",
    "    accs = {}\n",
    "    adj, features, labels = data.adj, data.features, data.labels\n",
    "    idx_train, idx_val, idx_test = data.idx_train, data.idx_val, data.idx_test\n",
    "    for i in range(retrain_iters):\n",
    "        if verbose:\n",
    "            print(\"\\nRetrain iter {}:\".format(i))\n",
    "        if model_type == 'GCNJaccard':\n",
    "            model = GCNJaccard(nfeat=features.shape[1], nclass=labels.max()+1,\n",
    "                               num_layers=num_layers,\n",
    "                               nhid=latent_size, device=device,\n",
    "                               weight_decay=weight_decay if weight_decay is not None else 5e-4,\n",
    "                               lr=lr if lr is not None else 0.01,\n",
    "                              )\n",
    "            model = model.to(device)\n",
    "            model.fit(features, adj, labels, idx_train, idx_val, train_iters=epochs, threshold=threshold, verbose=verbose)\n",
    "        elif model_type == 'RGCN':\n",
    "            model = RGCN(nnodes=adj.shape[0], nfeat=features.shape[1], nclass=labels.max()+1,\n",
    "                         num_layers=num_layers,\n",
    "                         nhid=latent_size, device=device,\n",
    "                         gamma=gamma if gamma is not None else 0.5,\n",
    "                         beta1=beta1 if beta1 is not None else 5e-4,\n",
    "                         beta2=weight_decay if weight_decay is not None else 5e-4,\n",
    "                        )\n",
    "            model = model.to(device)\n",
    "            model.fit(features, adj, labels, idx_train, idx_val, train_iters=epochs, verbose=verbose)\n",
    "        model.eval()\n",
    "        output = model.test(idx_test)\n",
    "        model.to(torch.device(\"cpu\"))\n",
    "        model.device = torch.device(\"cpu\")\n",
    "        models.append(model)\n",
    "        record_data(accs, output+[epochs], ['test_loss', 'test_acc', 'epochs'])\n",
    "\n",
    "    accs = {key + suffix: np.mean(value) for key, value in accs.items()}\n",
    "    return models, accs\n",
    "\n",
    "\n",
    "def train_multiple(data, loss_type, retrain_iters=5, suffix=\"\", verbose=True, device=None):\n",
    "    models = []\n",
    "    accs = {}\n",
    "    if struct_dropout_mode[0] == 'DNsampling' or (struct_dropout_mode[0] == 'standard' and len(struct_dropout_mode) == 3):\n",
    "        data.to(torch.device(\"cpu\"))\n",
    "        data_core = deepcopy(data)\n",
    "        add_distant_neighbors(data_core, struct_dropout_mode[-1])\n",
    "        data.to(device)\n",
    "        data_core.to(device)\n",
    "    else:\n",
    "        data_core = data.to(device)\n",
    "    for i in range(retrain_iters):\n",
    "        if verbose:\n",
    "            print(\"\\nRetrain iter {}:\".format(i))\n",
    "\n",
    "        model = GNN(\n",
    "            model_type=model_type,\n",
    "            num_features=info[\"num_features\"],\n",
    "            num_classes=info[\"num_classes\"],\n",
    "            reparam_mode=reparam_mode,\n",
    "            prior_mode=prior_mode,\n",
    "            struct_dropout_mode=struct_dropout_mode,\n",
    "            latent_size=latent_size,\n",
    "            num_layers=num_layers,\n",
    "            val_use_mean=val_use_mean,\n",
    "            reparam_all_layers=reparam_all_layers,\n",
    "            is_cuda=is_cuda,\n",
    "        )\n",
    "        record = train(\n",
    "            model,\n",
    "            data_core,\n",
    "            data_type=data_type,\n",
    "            model_type=model_type,\n",
    "            loss_type=loss_type,\n",
    "            epochs=epochs,\n",
    "            beta1_list=beta1_list,\n",
    "            beta2_list=beta2_list,\n",
    "            verbose=verbose,\n",
    "            isplot=False,\n",
    "            lr=lr,\n",
    "            weight_decay=weight_decay,\n",
    "            save_best_model=True,\n",
    "        )\n",
    "        record_data(accs, [record[\"train_acc\"][-1], record[\"val_acc\"][-1], record[\"test_acc\"][-1], record[\"epoch\"][-1], \n",
    "                           record[\"train_loss\"][-1], record[\"val_loss\"][-1], record[\"test_loss\"][-1],\n",
    "                           record[\"train_f1_micro\"][-1], record[\"best_val_f1_micro\"][-1], record[\"b_test_f1_micro\"][-1],\n",
    "                           record[\"train_f1_macro\"][-1], record[\"val_f1_macro\"][-1], record[\"test_f1_macro\"][-1]],     \n",
    "                    METRICS_RECORD)\n",
    "        model_best = load_model_dict_GNN(record[\"best_model_dict\"])\n",
    "        model.to_device(torch.device(\"cpu\"))\n",
    "        models.append((model, model_best))\n",
    "    accs = {key + suffix: np.mean(value) for key, value in accs.items()}\n",
    "    return models, accs\n",
    "\n",
    "\n",
    "def analyze(models, data, node_id, baseline, cached=None):\n",
    "    classification_margins = []\n",
    "    class_distrs = []\n",
    "    classification_margins_best = []\n",
    "    class_distrs_best = []\n",
    "    cpu_device = torch.device(\"cpu\")\n",
    "    if struct_dropout_mode[0] == 'DNsampling' or (struct_dropout_mode[0] == 'standard' and len(struct_dropout_mode) == 3):\n",
    "        data.to(cpu_device)\n",
    "        data_core = deepcopy(data)\n",
    "        add_distant_neighbors(data_core, struct_dropout_mode[-1])\n",
    "    else:\n",
    "        data_core = data.to(cpu_device)\n",
    "    for i, model in enumerate(models):\n",
    "        if isinstance(model, tuple):\n",
    "            model_last = model[0]\n",
    "            model_best = model[1]\n",
    "            model_best = model_best.to(cpu_device)\n",
    "            model_best.eval()\n",
    "        else:\n",
    "            model_last = model\n",
    "        model_last = model_last.to(cpu_device)\n",
    "        model_last.device = cpu_device\n",
    "        model_last.eval()\n",
    "        if baseline:\n",
    "            logits = model_last.predict(data_core.features, data_core.adj)\n",
    "            probs = torch.exp(logits[node_id])\n",
    "        else:\n",
    "            if cached is not None and hasattr(model_last, \"set_cache\"):\n",
    "                model_last.set_cache(cached)\n",
    "            logits, _ = model_last(data_core)\n",
    "            probs = F.softmax(logits[node_id], dim=-1)\n",
    "        class_distrs.append(probs)\n",
    "        best_second_class = (probs - 1000 * y_onehot[node_id]).argmax()\n",
    "        margin = probs[data_core.y[node_id]] - probs[best_second_class]\n",
    "        classification_margins.append(margin)\n",
    "\n",
    "        if isinstance(model, tuple):\n",
    "            if baseline:\n",
    "                logits_best = model_best.predict(data_core.features, data_core.adj)\n",
    "                probs_best = torch.exp(logits_best[node_id])\n",
    "            else:\n",
    "                if cached is not None and hasattr(model_best, \"set_cache\"):\n",
    "                    model_best.set_cache(cached)\n",
    "                logits_best, _ = model_best(data_core)\n",
    "                probs_best = F.softmax(logits_best[node_id], dim=-1)\n",
    "            class_distrs_best.append(probs_best)\n",
    "            best_second_class_best = (probs_best - 1000 * y_onehot[node_id]).argmax()\n",
    "            margin_best = probs_best[data_core.y[node_id]] - probs_best[best_second_class_best]\n",
    "            classification_margins_best.append(margin_best)\n",
    "\n",
    "    classification_margins = torch.stack(classification_margins)\n",
    "    class_distrs = torch.stack(class_distrs)\n",
    "    class_distrs, classification_margins = to_np_array(class_distrs, classification_margins)\n",
    "\n",
    "    if isinstance(model, tuple):\n",
    "        classification_margins_best = torch.stack(classification_margins_best)\n",
    "        class_distrs_best = torch.stack(class_distrs_best)\n",
    "        class_distrs_best, classification_margins_best = to_np_array(class_distrs_best, classification_margins_best)\n",
    "    return class_distrs, classification_margins, class_distrs_best, classification_margins_best\n",
    "\n",
    "\n",
    "# Visualization:\n",
    "def plot_attack(data, class_distrs_clean, class_distrs_attacked, node_id, retrain_iters):\n",
    "    def make_xlabel(ix, correct):\n",
    "        if ix==correct:\n",
    "            return \"Class {}\\n(correct)\".format(ix)\n",
    "        return \"Class {}\".format(ix)\n",
    "\n",
    "    figure = plt.figure(figsize=(12,4))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    center_ixs_clean = []\n",
    "    for ix, block in enumerate(class_distrs_clean.T):\n",
    "        x_ixs= np.arange(len(block)) + ix*(len(block)+2)\n",
    "        center_ixs_clean.append(np.mean(x_ixs))\n",
    "        color = '#555555'\n",
    "        if ix == data.y[node_id].item():\n",
    "            color = 'darkgreen'\n",
    "        plt.bar(x_ixs, block, color=color)\n",
    "\n",
    "    ax=plt.gca()\n",
    "    plt.ylim((-.05, 1.05))\n",
    "    plt.ylabel(\"Predicted probability\")\n",
    "    ax.set_xticks(center_ixs_clean)\n",
    "    ax.set_xticklabels([make_xlabel(k, data.y[node_id].item()) for k in range(info[\"num_classes\"])])\n",
    "    ax.set_title(\"Predicted class probabilities for node {} on clean data\\n({} re-trainings)\".format(node_id, retrain_iters))\n",
    "\n",
    "    fig = plt.subplot(1, 2, 2)\n",
    "    center_ixs_attacked = []\n",
    "    for ix, block in enumerate(class_distrs_attacked.T):\n",
    "        x_ixs= np.arange(len(block)) + ix*(len(block)+2)\n",
    "        center_ixs_attacked.append(np.mean(x_ixs))\n",
    "        color = '#555555'\n",
    "        if ix == data.y[node_id].item():\n",
    "            color = 'darkgreen'\n",
    "        plt.bar(x_ixs, block, color=color)\n",
    "\n",
    "\n",
    "    ax=plt.gca()\n",
    "    plt.ylim((-.05, 1.05))\n",
    "    ax.set_xticks(center_ixs_attacked)\n",
    "    ax.set_xticklabels([make_xlabel(k, data.y[node_id].item()) for k in range(info[\"num_classes\"])])\n",
    "    ax.set_title(\"Predicted class probabilities for node {} after {} perturbations\\n({} re-trainings)\".format(node_id, info_attack[\"n_perturbations\"], retrain_iters))\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Typical GIB-Cat setting: model_type=\"GAT\", beta1=0.001, beta2=0.01, struct_dropout_mode=(\"DNsampling\",'multi-categorical-sum',0.1,3,2);\n",
    "Typical GIB-Bern setting:model_type=\"GAT\", beta1=0.001, beta2=0.01, struct_dropout_mode=(\"DNsampling\",'Bernoulli',0.1,0.5,\"norm\",2);\n",
    "Standard GAT setting:    model_type=\"GAT\", beta1=-1,    beta2=-1,  struct_dropout_mode=(\"standard\", 0.6);\n",
    "Standard GCN setting:    model_type=\"GCN\", beta1=-1,    beta2=-1\n",
    "RGCN setting:            model_type=\"RGCN\"\n",
    "GCNJaccard setting:      model_type=\"GCNJaccard\"\n",
    "\"\"\"\n",
    "# Attack settings:\n",
    "attack_type = \"nettack\"     # Attack type. Currently only implemented Nettack.\n",
    "direct_attack = True        # Whether to use direct attack\n",
    "attacked_nodes = \"n:10\"     # Target nodes to attack. Use \"n:10\" for the default attack (40 nodes) in the Nettack paper. Alternatively, use \"r:1000\" for randomly choose 1000 test nodes to attack.\n",
    "n_perturbations = 1         # Number of feature or structural perturbations. Can be a integer (1, 2, 3,...). When set to -1, for each node will have degree + 2 perturbations.\n",
    "\n",
    "# Important model settings:\n",
    "exp_id = \"attack\"           # Experiment id, used for the directory name saving the experiment result files.\n",
    "data_type = 'Cora-bool'     # Data type. Should use suffix \"-bool\" for binarized features. Choose from \"Cora-bool\", \"Pubmed-bool\" and \"citeseer-bool\"\n",
    "model_type = 'GAT'          # Name of the base model. Choose from \"GAT\", \"GCN\", 'GCNJaccard', 'RGCN'. \n",
    "                            # For GIB-Cat and GIB-Bern, still choose model_type=\"GAT\", but set either beta1 or beta2 nonzero.\n",
    "beta1 = 0.001               # coefficient for the XIB term. If -1, this term will turn off.\n",
    "beta2 = 0.01                # coefficient for the AIB term. If -1, this term will have 0 coefficent (but may still perform sampling, depending on \"struct_dropout_mode\")\n",
    "struct_dropout_mode = (\"DNsampling\",'multi-categorical-sum',1,3,2)  # Mode for how the structural representation is generated.\n",
    "                            # For GIB-Cat, choose from (\"Nsampling\", 'multi-categorical-sum', 0.1, 3) (here 0.1 is temperature, k=3 is the number of sampled neighboring edges with replacement),\n",
    "                            #    and (\"DNsampling\", 'multi-categorical-sum', 0.1, 3, 2) (similar as above, but with 2-hop neighbors)\n",
    "                            # For GIB-Bern, choose from (\"Nsampling\",'Bernoulli',0.1,0.5,\"norm\") (here 0.1 is temperature, 0.5 is the prior for the Bernoulli probability)\n",
    "                            #    and (\"DNsampling\",'Bernoulli',0.1,0.5,\"norm\",2) (with 2-hop neighbors)\n",
    "                            # For standard GAT, choose from (\"standard\", 0.6) (where standard dropout used on the attention weights in GAT)\n",
    "                            #    and (\"standard\", 0.6, 2) (with 2-hop neighbors)\n",
    "latent_size = 16            # Latent dimension for GCN-based or GAT-based models. \n",
    "sample_size = 1             # How many Z sampled from each node X.\n",
    "num_layers = 2              # Number of layers for the GNN.\n",
    "reparam_mode = \"diag\"       # Reparameterization mode for XIB. Choose from \"None\", \"diag\" or \"full\"\n",
    "prior_mode = \"mixGau-100\"   # Prior mode. Choose from \"Gaussian\" or \"mixGau-100\" (mixture of 100 Gaussian components)\n",
    "is_anneal_beta = True       # Whether to anneal beta1 and beta2 from 0 up during training. Default True.\n",
    "val_use_mean = True         # Whether during evaluation use the parameter value instead of sampling. If True, during evaluation,\n",
    "                            # XIB will use mean for prediction, and AIB will use the parameter of the categorical distribution for prediction.\n",
    "reparam_all_layers = (-2,)  # Which layers to use XIB, e.g. (1,2,4). Default (-2,), meaning the second last layer. If True, use XIB for all layers.\n",
    "epochs = 2000               # Number of epochs. Default 2000\n",
    "lr = -1                     # Learning rate. If -1, use default learning rate for each model\n",
    "weight_decay = -1           # weight decay. If -1, use default weight decay for each model\n",
    "date_time = \"{0}-{1}\".format(datetime.datetime.now().month, datetime.datetime.now().day)  # Today's month and day. Used for the directory name saving the experiment result files.\n",
    "seed = 0                    # Random seed.\n",
    "idx = \"0\"                   # idx to differentiate different files. Only used if wanting to run the same setting for many times.\n",
    "\n",
    "# Other settings:\n",
    "retrain_iters = 1           # How many models to train for each setting and given attacked nodes.\n",
    "is_load = True              # Whether to load previous checkpoint.\n",
    "threshold = 0.05            # Threshold for GCNJaccard.\n",
    "gamma = 0.5                 # gamma for RGCN\n",
    "\n",
    "try:\n",
    "    get_ipython().run_line_magic('matplotlib', 'inline')\n",
    "    import matplotlib.pylab as plt\n",
    "    isplot = False\n",
    "    verbose = False\n",
    "except:\n",
    "    import matplotlib\n",
    "    isplot = False\n",
    "    verbose = False\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--exp_id', default=exp_id, help='experiment ID')\n",
    "    parser.add_argument('--data_type', help='Data type: choose from PROTEINS.', required=True)\n",
    "    parser.add_argument('--model_type', default=\"GAT\", help='Model type: GCN or GAT or GCNJaccard or RGCN')\n",
    "    parser.add_argument('--direct_attack', type=str2bool, nargs='?', const=True, default=True, help='Whether to do direct attack.')\n",
    "    parser.add_argument('--attacked_nodes', default=\"::\", help='indices of attacked_nodes')\n",
    "    parser.add_argument('--n_perturbations', type=int, default=5, help='number of perturbations')\n",
    "    parser.add_argument('--beta1', type=float, default=0.001, help='beta1 value for feature IB, set a float value >= 0.')\n",
    "    parser.add_argument('--beta2', type=float, default=0.01, help='beta2 value for structure IB, set a float value >= 0.')\n",
    "    parser.add_argument('--latent_size', type=int, default=16, help='latent_size')\n",
    "    parser.add_argument('--sample_size', type=int, default=1, help='sample_size')\n",
    "    parser.add_argument('--num_layers', type=int, default=2, help='num_layers')\n",
    "    parser.add_argument('--reparam_mode', default=\"diag\", help='diag, diagg, or full')\n",
    "    parser.add_argument('--prior_mode', default=\"Gaussian\", help='prior mode for VIB')\n",
    "    parser.add_argument('--struct_dropout_mode', default=\"(standard, 0.6)\", help='mode for structure dropout.')\n",
    "    parser.add_argument('--is_anneal_beta', type=str2bool, nargs='?', const=True, default=True, help='Whether to anneal beta.')\n",
    "    parser.add_argument('--val_use_mean', type=str2bool, nargs='?', const=True, default=True, help='Whether to use mean of Z during validation.')\n",
    "    parser.add_argument('--reparam_all_layers', help='Whether to reparameterize all layers.')\n",
    "    parser.add_argument('--epochs', type=int, default=1000, help=\"Number of epochs.\")\n",
    "    parser.add_argument('--date_time', default=date_time, help=\"Current date and time.\")\n",
    "    parser.add_argument('--lr', type=float, default=lr, help='lr')\n",
    "    parser.add_argument('--weight_decay', type=float, default=-1, help=\"weight_decay.\")\n",
    "    parser.add_argument('--seed', type=int, default=1, help='seed')\n",
    "    parser.add_argument('--gpuid', help='an integer for the accumulator', required=True)\n",
    "    parser.add_argument('--idx', default=\"0\", help='idx')\n",
    "    parser.add_argument('--threshold', type=float, default=0.05, help='threshold for GCNJaccard')\n",
    "    parser.add_argument('--gamma', type=float, default=0.3, help='gamma for RGCN')\n",
    "    parser.add_argument('--retrain_iters', type=int, default=1, help='retrain_iters')\n",
    "    parser.add_argument('--is_load', type=str2bool, nargs='?', const=True, default=True, help='Whether to load previous trained instance.')\n",
    "    \n",
    "    args = parser.parse_args()\n",
    "\n",
    "\n",
    "if \"args\" in locals():\n",
    "    exp_id = args.exp_id\n",
    "    data_type = args.data_type\n",
    "    model_type = args.model_type\n",
    "    direct_attack = args.direct_attack\n",
    "    attacked_nodes = args.attacked_nodes\n",
    "    n_perturbations = args.n_perturbations\n",
    "    beta1 = args.beta1\n",
    "    beta2 = args.beta2\n",
    "    latent_size = args.latent_size\n",
    "    sample_size = args.sample_size\n",
    "    num_layers = args.num_layers\n",
    "    reparam_mode = args.reparam_mode\n",
    "    prior_mode = args.prior_mode\n",
    "    struct_dropout_mode = eval_tuple(args.struct_dropout_mode)\n",
    "    is_anneal_beta = args.is_anneal_beta\n",
    "    val_use_mean = args.val_use_mean\n",
    "    reparam_all_layers = eval_tuple(args.reparam_all_layers)\n",
    "    epochs = args.epochs\n",
    "    date_time = args.date_time\n",
    "    lr = args.lr\n",
    "    weight_decay = args.weight_decay\n",
    "    seed = args.seed\n",
    "    idx = args.idx\n",
    "    threshold = args.threshold\n",
    "    gamma = args.gamma\n",
    "    retrain_iters = args.retrain_iters\n",
    "    is_load = args.is_load\n",
    "    is_cuda = eval(args.gpuid)\n",
    "    if not isinstance(is_cuda, bool):\n",
    "        is_cuda = \"cuda:{}\".format(is_cuda)\n",
    "baseline = model_type in ['GCNJaccard', 'RGCN']\n",
    "\n",
    "\n",
    "device = torch.device(is_cuda if isinstance(is_cuda, str) else \"cuda\" if is_cuda else \"cpu\")\n",
    "\n",
    "dirname = GIB_PATH + \"/{0}_{1}/\".format(exp_id, date_time)\n",
    "if baseline: ########## difference is that for baselines we have an additional parameters\n",
    "    filename = dirname + \"{0}_{1}_dire_{2}_nodes_{3}_pert_{4}_lat_{5}_samp_{6}_l_{7}_anl_{8}_mean_{9}_reall_{10}_epochs_{11}_lr_{12}_l2_{13}_seed_{14}_threshold_{15}_gamma_{16}_{17}_id_{18}\".format(\n",
    "        data_type, model_type, direct_attack, attacked_nodes, n_perturbations,\n",
    "        latent_size, sample_size, num_layers, \n",
    "        is_anneal_beta, val_use_mean, reparam_all_layers, epochs, lr, weight_decay, seed, threshold, gamma, is_cuda, idx,\n",
    "    )\n",
    "else:\n",
    "    filename = dirname + \"{0}_{1}_dire_{2}_nodes_{3}_pert_{4}_beta_{5}_{6}_lat_{7}_samp_{8}_l_{9}_reparam_{10}_prior_{11}_sdrop_{12}_anl_{13}_mean_{14}_reall_{15}_epochs_{16}_lr_{17}_l2_{18}_seed_{19}_{20}_id_{21}\".format(\n",
    "        data_type, model_type, direct_attack, attacked_nodes, n_perturbations, beta1, beta2,\n",
    "        latent_size, sample_size, num_layers, reparam_mode, prior_mode, to_string(struct_dropout_mode, \"-\"), \n",
    "        is_anneal_beta, val_use_mean, reparam_all_layers, epochs, lr, weight_decay, seed, is_cuda, idx,\n",
    "    )\n",
    "make_dir(filename)\n",
    "print(filename)\n",
    "\n",
    "np.random.rand(seed)\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "# Set up attacked nodes:\n",
    "assert attacked_nodes in [\"r:1000\", \"n:10\"]\n",
    "if attacked_nodes.startswith(\"n:\"):\n",
    "    data, info = get_data(data_type, seed=seed)\n",
    "else:\n",
    "    data, info = get_data(data_type, seed=seed, attacked_nodes=attacked_nodes)\n",
    "data = process_data_for_nettack(data).to(device)\n",
    "if attacked_nodes.startswith(\"n:\"):\n",
    "    node_ids = select_nodes(data, info, data_type, num=eval(attacked_nodes.split(\":\")[-1]), seed=seed)\n",
    "else:\n",
    "    node_ids = info['node_ids']\n",
    "y_onehot = torch.eye(info[\"num_classes\"])[data.y]\n",
    "\n",
    "\n",
    "if n_perturbations == -1:\n",
    "    n_perturbations = None\n",
    "if struct_dropout_mode[0] is None:\n",
    "    struct_dropout_mode = (\"None\",)\n",
    "if lr == -1:\n",
    "    lr = None\n",
    "if weight_decay == -1:\n",
    "    weight_decay = None\n",
    "\n",
    "if beta1 == -1:\n",
    "    beta1 = None\n",
    "if beta1 is None:\n",
    "    beta1_list, reparam_mode, prior_mode = None, None, None\n",
    "else:\n",
    "    if is_anneal_beta:\n",
    "        beta_init = 0\n",
    "        init_length = int(epochs / 4)\n",
    "        anneal_length = int(epochs / 4)\n",
    "        beta_inter = Beta_Function(np.linspace(0,1,anneal_length),1,4)\n",
    "        beta1_inter = beta_inter / 4 * (beta_init - beta1) + beta1\n",
    "        beta1_list = np.concatenate([np.ones(init_length) * beta_init, beta1_inter, \n",
    "                                     np.ones(epochs - init_length - anneal_length + 1) * beta1])\n",
    "    else:\n",
    "        beta1_list = np.ones(epochs + 1) * beta1\n",
    "if beta2 == -1:\n",
    "    beta2_list = None\n",
    "else:\n",
    "    if is_anneal_beta:\n",
    "        beta_init = 0\n",
    "        init_length = int(epochs / 4)\n",
    "        anneal_length = int(epochs / 4)\n",
    "        beta_inter = Beta_Function(np.linspace(0,1,anneal_length),1,4)\n",
    "        beta2_inter = beta_inter / 4 * (beta_init - beta2) + beta2\n",
    "        beta2_list = np.concatenate([np.ones(init_length) * beta_init, beta2_inter, \n",
    "                                     np.ones(epochs - init_length - anneal_length + 1) * beta2])\n",
    "    else:\n",
    "        beta2_list = np.ones(epochs + 1) * beta2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Surrogate model for Nettack:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if attack_type == \"nettack\":\n",
    "    surrogate_model = GCN(nfeat=info[\"num_features\"], nclass=info[\"num_classes\"],\n",
    "                nhid=16, dropout=0, with_relu=False, with_bias=False, device=device).to(device)\n",
    "    surrogate_model.fit(data.features, data.adj, data.labels, data.idx_train)\n",
    "    W1 = to_np_array(surrogate_model.gc1.weight)\n",
    "    W2 = to_np_array(surrogate_model.gc2.weight)\n",
    "    all_dict = {\"params\": {\"W1\": W1, \"W2\": W2}}\n",
    "else:\n",
    "    raise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform and evaluate attacks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load previous instance:\n",
    "if is_load:\n",
    "    try:\n",
    "        filename_core = \"_\".join(filename.split(\"/\")[-1].split(\"_\")[:-3])\n",
    "        cand_filename = filter_filename(dirname, include=filename_core)\n",
    "        if len(cand_filename) == 0:\n",
    "            raise Exception(\"Did not find previous file. Create new.\")\n",
    "        assert cand_filename[0].endswith(\".p\")\n",
    "        filename = dirname + cand_filename[0][:-2]\n",
    "        all_dict_cand = pickle.load(open(filename + \".p\", \"rb\"))\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        is_load = False\n",
    "    if is_load:\n",
    "        if attack_type == \"nettack\":\n",
    "            if not \"params\" in all_dict_cand:\n",
    "                is_load = False\n",
    "            else:\n",
    "                all_dict = all_dict_cand\n",
    "                surrogate_model.gc1.weight.data = torch.FloatTensor(all_dict[\"params\"][\"W1\"]).to(device)\n",
    "                surrogate_model.gc2.weight.data = torch.FloatTensor(all_dict[\"params\"][\"W2\"]).to(device)\n",
    "                print(\"Load previous trained instance at {}.\".format(filename))\n",
    "\n",
    "is_train = True\n",
    "if is_load:\n",
    "    is_all_attacked = True\n",
    "    for node_id in node_ids:\n",
    "        if node_id not in all_dict:\n",
    "            is_all_attacked = False\n",
    "            break\n",
    "    if not is_all_attacked:\n",
    "        if \"models_before\" in all_dict and isinstance(all_dict[\"models_before\"][0], tuple):\n",
    "            print(\"Load models_before.\")\n",
    "            models_before = [(load_model_dict_GNN(model_dict_last, is_cuda=is_cuda), \n",
    "                              load_model_dict_GNN(model_dict_best, is_cuda=is_cuda)) for model_dict_last, model_dict_best in all_dict[\"models_before\"]]\n",
    "            is_train = False\n",
    "    else:\n",
    "        is_train = False\n",
    "        print(\"all node_ids are attacked. Skip training clean model.\")\n",
    "\n",
    "if is_train:\n",
    "    print(\"Training clean model:\")\n",
    "    if baseline:\n",
    "        models_before, accs_clean = train_baseline_multiple(data, model_type, device, retrain_iters=retrain_iters, suffix=\"_clean\", verbose=verbose)\n",
    "    else:\n",
    "        models_before, accs_clean = train_multiple(data, loss_type=info[\"loss\"], retrain_iters=retrain_iters, suffix=\"_clean\", verbose=verbose, device=device)\n",
    "        all_dict[\"models_before\"] = [(model_last.model_dict, model_best.model_dict) for model_last, model_best in models_before]\n",
    "else:\n",
    "    accs_clean = {}\n",
    "    for key in METRICS_RECORD:\n",
    "        accs_clean[key + \"_clean\"] = all_dict[node_ids[0]][key + \"_clean\"]\n",
    "\n",
    "# Iterations:\n",
    "for i, node_id in enumerate(node_ids):\n",
    "    node_id = int(node_id)\n",
    "    if is_load and node_id in all_dict:\n",
    "        print(\"Skip node {}.\".format(node_id))\n",
    "        continue\n",
    "    print(\"\\nAttacking the {}th node id={}:\".format(i, node_id))\n",
    "    if attack_type == \"nettack\":\n",
    "        data_attacked, info_attack = get_attacked_data_deeprobust(\n",
    "            data,\n",
    "            surrogate_model=surrogate_model,\n",
    "            target_node=node_id,\n",
    "            direct_attack=direct_attack,\n",
    "            n_perturbations=n_perturbations,\n",
    "            verbose=False,\n",
    "        )\n",
    "    else:\n",
    "        raise\n",
    "\n",
    "    # Train with clean data:\n",
    "    if verbose:\n",
    "        print(\"\\nClean data:\")\n",
    "    class_distrs_clean, classification_margins_clean, class_distrs_clean_best, classification_margins_clean_best = analyze(models_before, data, node_id, baseline)\n",
    "    class_distrs_evasive, classification_margins_evasive, class_distrs_evasive_best, classification_margins_evasive_best = analyze(models_before, data_attacked, node_id, baseline, cached=False)\n",
    "\n",
    "    # Train with attacked data:\n",
    "    if verbose:\n",
    "        print(\"\\nAttacked data:\")\n",
    "    if baseline:\n",
    "        models_attacked, accs_attacked = train_baseline_multiple(data_attacked, model_type, device, retrain_iters=retrain_iters, suffix=\"_attacked\", verbose=verbose)\n",
    "    else:\n",
    "        models_attacked, accs_attacked = train_multiple(data_attacked, loss_type=info[\"loss\"], retrain_iters=retrain_iters, suffix=\"_attacked\", verbose=verbose, device=device)\n",
    "    class_distrs_attacked, classification_margins_attacked, class_distrs_attacked_best, classification_margins_attacked_best = analyze(models_attacked, data_attacked, node_id, baseline)\n",
    "\n",
    "    all_dict[node_id] = {\n",
    "        \"class_distrs_clean\": class_distrs_clean,\n",
    "        \"class_distrs_evasive\": class_distrs_evasive,\n",
    "        \"class_distrs_attacked\": class_distrs_attacked,\n",
    "        \"classification_margins_clean\": classification_margins_clean,\n",
    "        \"classification_margins_evasive\": classification_margins_evasive,\n",
    "        \"classification_margins_attacked\": classification_margins_attacked,  \n",
    "        \"class_distrs_clean_best\": class_distrs_clean_best,\n",
    "        \"class_distrs_evasive_best\": class_distrs_evasive_best,\n",
    "        \"class_distrs_attacked_best\": class_distrs_attacked_best,\n",
    "        \"classification_margins_clean_best\": classification_margins_clean_best,\n",
    "        \"classification_margins_evasive_best\": classification_margins_evasive_best,\n",
    "        \"classification_margins_attacked_best\": classification_margins_attacked_best,  \n",
    "    }\n",
    "    all_dict[node_id].update(accs_clean)\n",
    "    all_dict[node_id].update(accs_attacked)\n",
    "    if len(class_distrs_attacked_best) > 0:\n",
    "        all_dict[node_id][\"models_before\"] = [(model_last.model_dict, model_best.model_dict) for model_last, model_best in models_before]\n",
    "        all_dict[node_id][\"models_attacked\"] = [(model_last.model_dict, model_best.model_dict) for model_last, model_best in models_attacked]\n",
    "    pickle.dump(all_dict, open(filename + \".p\", \"wb\"))\n",
    "    print(\"margins: clean: {:.4f}   evasive: {:.4f}   poison: {:.4f}   test_acc_clean: {:.4f}   test_acc_poison: {:.4f}     epochs: ({:.0f}, {:.0f})\".format(\n",
    "        classification_margins_clean.mean(),\n",
    "        classification_margins_evasive.mean(),\n",
    "        classification_margins_attacked.mean(),\n",
    "        all_dict[node_id][\"test_acc_clean\"],\n",
    "        all_dict[node_id][\"test_acc_attacked\"],\n",
    "        all_dict[node_id][\"epochs_clean\"],\n",
    "        all_dict[node_id][\"epochs_attacked\"],\n",
    "    ))\n",
    "    try:\n",
    "        sys.stdout.flush()\n",
    "    except:\n",
    "        pass\n",
    "    if isplot:\n",
    "        plot_attack(data, class_distrs_clean, class_distrs_attacked, node_id, retrain_iters)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python base",
   "language": "python",
   "name": "base"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
